Skip to content

Add Python parser via tree-sitter - #23

Open
awe-srush wants to merge 1 commit into
google:mainfrom
awe-srush:feat/python-parser
Open

Add Python parser via tree-sitter#23
awe-srush wants to merge 1 commit into
google:mainfrom
awe-srush:feat/python-parser

Conversation

@awe-srush

Copy link
Copy Markdown

Summary

Adds PythonParser, the first concrete implementation of TreeSitterParserBase, enabling Vanir to detect missing security patches in Python source code. Implements all three language-specific abstract methods required by the base class.

Depends on:

  • feat/tree-sitter-integration (base class)
  • bug/fix-subclass-discovery (recursive subclass registration)
  • feat/python-parser-deps (tree-sitter pip dependencies)
  • feat/python-parser-build (Bazel build targets)

Objective

Extend Vanir's patch detection pipeline to support .py files by plugging a tree-sitter-backed Python parser into the existing AbstractLanguageParser contract. Once registered, the parser is auto-discovered by get_parser_class() and requires no changes to any upstream Vanir logic.

Files

File Role
vanir/language_parsers/python/python_parser.py PythonParser implementation
vanir/language_parsers/python/__init__.py Package marker
vanir/language_parsers/language_parsers.py Adds python import to register the parser

Implementation

Module-level setup

Tree-sitter imports are guarded with try/except ImportError so the module is safely importable on Python 3.9, where tree-sitter and tree-sitter-python are not installed:

try:
    import tree_sitter_python as _tspython
    from tree_sitter import Language, QueryCursor
    _PY_LANGUAGE = Language(_tspython.language())
    _FUNC_QUERY  = _PY_LANGUAGE.query(...)
    _LOCALS_QUERY = _PY_LANGUAGE.query(...)
    _TREE_SITTER_AVAILABLE = True
except ImportError:
    _TREE_SITTER_AVAILABLE = False

Two queries are compiled once at module import (never per file):

  • _FUNC_QUERY, matches every function_definition node, capturing:
    • @func.def, the full function node
    • @func.name, the function name identifier
    • @func.body, the function body block
  • _LOCALS_QUERY, matches five assignment/call patterns to collect local variables and called functions:
    • (assignment left: (_) @assign.lhs), standard assignments
    • (augmented_assignment left: (identifier) @aug.lhs), +=, -= etc.
    • (for_statement left: (_) @for.lhs), loop variables
    • (named_expression name: (identifier) @walrus.name), walrus operator :=
    • (call function: (identifier) @call.name), direct function calls

PythonParser (class)

Thin subclass of TreeSitterParserBase. Sets four class-level attributes:

  • LANGUAGE = _PY_LANGUAGE, compiled Python grammar
  • _FUNC_QUERY = _FUNC_QUERY, function discovery query
  • SKIP_TOKEN_TYPES, prunes comment, newline, indent, dedent, line_continuation from token collection
  • STRING_NODE_TYPE = 'string', Python string nodes emitted as single tokens

get_supported_extensions()

Returns ['.py'] when tree-sitter is available, [] on Python 3.9, causing the dispatcher to silently skip .py files rather than raise an error.

Abstract method implementations

_extract_param_names(params_node) → list[str]

Walks params_node.named_children and resolves each to a plain name string via the module-level helper _get_param_name().

_get_param_name() handles all Python parameter node types:

  • identifier, plain param (x)
  • typed_parameter, annotated param (x: int)
  • default_parameter, param with default (x=0)
  • typed_default_parameter, annotated with default (x: int = 0)
  • list_splat_pattern, variadic positional (*args)
  • dictionary_splat_pattern, variadic keyword (**kwargs)
  • bare * separator, returned as None and skipped

For all compound types, the function finds the first identifier child node and decodes its text.

_extract_annotations(func_node) → (return_types, used_data_types)

Extracts Python type annotations from a function_definition node:

Return type annotation:

  • Retrieves the return_type field via child_by_field_name('return_type')
  • Flattens tokens from the annotation node using _flat_tokens_cursor()
  • Result: return_types = [[token, ...]] or [] if no annotation

Parameter type annotations:

  • Iterates parameters named children
  • For nodes of type typed_parameter or typed_default_parameter, retrieves the type field via child_by_field_name('type')
  • Flattens tokens for each annotation
  • Result: used_data_types = [[token, ...], ...], one list per annotated param

_collect_locals_calls(body_node, nested_ranges) → (local_vars, called_fns)

Uses QueryCursor(_LOCALS_QUERY).captures(body_node) to collect all assignment targets and call names within the function body.

Each captured node is filtered through _is_nested(), any node whose start_byte falls within a nested function's byte range is excluded, preventing inner-scope variables from leaking into the outer function's Vanir signature.

Capture groups processed:

  • assign.lhs, passed to _collect_ids_from_lhs() which handles plain identifiers, tuple unpacking (a, b = ...), and pattern lists
  • aug.lhs, augmented assignment targets, added directly if identifier
  • for.lhs, for-loop variables, passed to _collect_ids_from_lhs()
  • walrus.name, walrus operator targets (:=), added if identifier
  • call.name, direct function call names, added to called_fns

Runtime availability

Vanir runs on Python 3.9 through 3.13. The tree-sitter packages are not installed for Python 3.9, so when Vanir itself runs on 3.9, get_supported_extensions() returns [] and .py files are silently skipped with no error.

On Python 3.10 and above, the parser is fully active.

Debug harness

python_parser.py includes a __main__ block for local inspection during development. It is never executed when Vanir imports the module.

Run directly to inspect parser output on any .py file:

python vanir/language_parsers/python/python_parser.py path/to/file.py

If no file is provided, it runs against a built-in sample covering classes, annotated methods, async functions, walrus operator, and tuple unpacking.

Prints all extracted chunks (name, parameters, return types, local variables, called functions, tokens, normalized form) and also runs an error handling test on intentionally broken Python to verify ParseError collection.

Dependencies

Must be merged after:

  • feat/tree-sitter-integration
  • bug/fix-subclass-discovery
  • feat/python-parser-deps
  • feat/python-parser-build

docs/python-parser follows this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant